# -*- coding: utf-8 -*-
# @Author: Tim Liu
# @Date: 2025-01-13
# @Last Modified by: Cursor
# @Last Modified time: 2025-01-14
# @Description: API endpoints for graph index management

from fastapi import APIRouter, Depends, Security
from sqlalchemy.ext.asyncio import AsyncSession

from config.settings import *
from utils.response import SuccessResponse
from core.database import db_getter
from core.exception import CustomException

from crewplus.helper.verify import get_api_key
from crewplus.apps.apikey.models.account_api_key import AccountApiKey
from crewplus.apps.rag.graph.knowledge_graph_builder import KnowledgeGraphBuilder
from crewplus.apps.rag.schemas.graph_index import GraphIndexRequest

from crewplus.services.vdb_service import VDBService
from crewplus.services.graph_engines.neo4j_engine import Neo4jEngine

app = APIRouter()

@app.post("/graph/flush_index", summary="Flush graph index")
async def flush_graph_index(request: GraphIndexRequest, apikey: AccountApiKey = Security(get_api_key), db: AsyncSession = Depends(db_getter)):
    """Flush graph index for the given API key and optional schema content."""
    if apikey is None:
        raise CustomException("Invalid API key", code=400)

    try:
        neo4j_engine = Neo4jEngine(
            url=NEO4J_GRAPH_ENDPOINT,
            user=NEO4J_USER,
            password=NEO4J_PASSWORD
        )

        # Create KnowledgeGraphBuilder instance
        graph_builder = KnowledgeGraphBuilder(db=db)
        graph_builder.register_graph_instance(user_id=apikey.user_id, graph_instance=neo4j_engine)
        
        # Bind embedder and vector store
        vdbsrv = VDBService()
        embeddings = vdbsrv.get_embeddings()
        graph_builder.bind_embedder(embeddings)
            
        graph_builder.flush_vector_index(user_id=apikey.user_id, schema_content=request.schema_content)
    except Exception as e:
        raise CustomException(f"Error flushing graph index: {str(e)}", code=400)
    
    return SuccessResponse("Graph index flushed successfully.")

@app.post("/graph/upsert_embeddings", summary="Upsert graph embeddings")
async def upsert_graph_embeddings(request: GraphIndexRequest, apikey: AccountApiKey = Security(get_api_key), db: AsyncSession = Depends(db_getter)):
    """Upsert graph embeddings for the given API key and optional schema content."""
    if apikey is None:
        raise CustomException("Invalid API key", code=400)

    try:
        neo4j_engine = Neo4jEngine(
            url=NEO4J_GRAPH_ENDPOINT,
            user=NEO4J_USER,
            password=NEO4J_PASSWORD
        )

        # Create KnowledgeGraphBuilder instance
        graph_builder = KnowledgeGraphBuilder(db=db)
        graph_builder.register_graph_instance(user_id=apikey.user_id, graph_instance=neo4j_engine)
        
        # Bind embedder and vector store
        vdbsrv = VDBService()
        embeddings = vdbsrv.get_embeddings()
        graph_builder.bind_embedder(embeddings)
                    
        graph_builder.upsert_embeddings_for_nodes(user_id=apikey.user_id, schema_content=request.schema_content)
    except Exception as e:
        raise CustomException(f"Error upserting graph embeddings: {str(e)}", code=400)
    
    return SuccessResponse("Graph embeddings upserted successfully.")